import sys, os, glob
import numpy as np
from astropy.io import fits

# 参数检查
if len(sys.argv) != 4:
    print("HELP: python filter_cands_simple.py /path/to/cand_output /path/to/filtered.txt /path/to/fits_dir")
    sys.exit(1)

cand_dir = sys.argv[1]
output_txt = sys.argv[2]
fits_dir = sys.argv[3]

DM_min = 1000
DM_max = 1300
threshold_snr = 7
min_members = 2
max_filter = 10

# 获取 cand 文件
cand_files = sorted(glob.glob(os.path.join(cand_dir, "*.cand")))
print(f"📂 Found {len(cand_files)} cand files.")

# 获取 .fits 文件，计算样本数（只使用第一个）
fits_files = sorted(glob.glob(os.path.join(fits_dir, "*.fits")))
if not fits_files:
    raise RuntimeError("未找到任何 .fits 文件")
with fits.open(fits_files[0]) as hdul:
    h = hdul[1].header
    samples_per_fits = h['NSBLK'] * h['NAXIS2']

if os.path.exists(output_txt):
    with open(output_txt, 'r') as f:
        processed_lines = f.read()
else:
    processed_lines = ''

# 筛选并恢复每条候选体所在的相对 sample
with open(output_txt, 'a+') as f_out:
    for cand_file in cand_files:
        if os.path.getsize(cand_file) == 0:
            continue

        data = np.loadtxt(cand_file)
        fits_name = os.path.basename(cand_file).replace(".cand", ".fits")

        if data.ndim == 1:
            data = np.expand_dims(data, axis=0)

        for row in data:
            snr = row[0]
            start_samp = int(row[1])
            N_filter = int(row[3])
            DM = float(row[5])
            N_member = int(row[6])

            # 获取候选体在合并数据中的 sample，然后还原回原始 fits 的 index 和相对 sample
            file_index = start_samp // samples_per_fits
            local_sample = start_samp % samples_per_fits

            if file_index >= len(fits_files):
                continue  # 超出范围跳过

            fits_name = os.path.basename(fits_files[file_index])
            tag = f"{fits_name} {local_sample}"

            # if (DM_min < DM < DM_max and snr > threshold_snr and
            #     N_member > min_members and N_filter < max_filter and
            #     tag not in processed_lines):
                
            #     f_out.write(f"{fits_name} {local_sample} {DM:.2f} {N_member} {N_filter} {snr:.2f}\n")
            if snr > 3 and tag not in processed_lines:
                f_out.write(f"{fits_name} {local_sample} {DM:.2f} {N_member} {N_filter} {snr:.2f}\n")

print(f"======筛选完成，已保存至 {output_txt}")